import json
import random
import logging
from typing import Dict, List, Optional, Any

# 配置日志记录器
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('GenerateReplyPrompt')

class GenerateReplyPrompt:
    """
    用于生成符合患者人格特征的回复提示词的工具类
    """
    
    # 常量定义
    DEFAULT_PERSONALITY_FILE = "/mnt/ssd/jsj/patient/script/人格定义.json"
    MEMORY_CONSTRAINT_KEYWORDS = ["哪不舒服", "症状", "哪里痛", "怎么了", "哪里不舒服", "感觉如何"]
    
    def __init__(self, personality_definition_file: str = DEFAULT_PERSONALITY_FILE):
        """
        初始化提示词生成器，加载人格定义
        
        Args:
            personality_definition_file: 人格定义JSON文件路径
        """
        self.personality_definition_file = personality_definition_file
        self.personality_definition = self.load_personality_definition(personality_definition_file)
    
    def load_personality_definition(self, file_path: str) -> Dict[str, Any]:
        """
        加载人格定义文件
        
        Args:
            file_path: 文件路径
            
        Returns:
            人格定义字典
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
                return json.loads(content)
        except FileNotFoundError:
            logger.error(f"人格定义文件未找到: {file_path}")
            # 返回默认的空结构
            return {"personality_dimensions": {}}
        except json.JSONDecodeError as e:
            logger.error(f"人格定义文件格式错误: {file_path}, 错误: {e}")
            return {"personality_dimensions": {}}
        except Exception as e:
            logger.error(f"加载人格定义文件失败: {file_path}, 错误: {e}")
            return {"personality_dimensions": {}}
    
    def get_personality_description(self, target_personality: Dict[str, str]) -> str:
        """
        根据目标人格生成详细的人格描述文本
        
        Args:
            target_personality: 目标人格字典
            
        Returns:
            人格描述文本
        """
        if not target_personality:
            logger.warning("目标人格为空，返回默认描述")
            return "患者人格特征详细定义：暂无特定人格信息"
        
        descriptions = []
        descriptions.append("【患者人格特征详细定义】：")
        
        # 直接遍历目标人格中的所有键值对，确保包含所有人格维度
        for dim_name, current_value in target_personality.items():
            # 跳过认知状态，因为在定义文件中标记为暂不考虑
            if dim_name == "认知状态":
                continue
            
            # 对于所有人格维度，都从定义文件中获取信息
            dim_info = None
            for _, info in self.personality_definition.get("personality_dimensions", {}).items():
                if info.get("name") == dim_name:
                    dim_info = info
                    break
            
            if dim_info:
                self._append_detailed_dimension_info(descriptions, dim_name, current_value, dim_info)
            else:
                # 如果在定义文件中找不到对应的维度信息，仍然添加基本信息
                descriptions.append(f"{dim_name}: {current_value}")
        
        return '\n'.join(descriptions)
    
    def _append_detailed_dimension_info(self, descriptions: List[str], dim_name: str, 
                                       current_value: str, dim_info: Dict[str, Any]) -> None:
        """
        添加详细的维度信息到描述列表中
        
        Args:
            descriptions: 描述列表
            dim_name: 维度名称
            current_value: 当前值
            dim_info: 维度信息
        """
        dim_key_name = dim_info.get("key", dim_name)
        values_info = dim_info.get("values", {})
        value_info = values_info.get(current_value, {})
        
        desc_parts = [f"{dim_name}（{dim_key_name}）: {current_value}"]
        
        if "定义" in value_info:
            desc_parts.append(f"  - 定义：{value_info['定义']}")
        if "表现" in value_info:
            desc_parts.append(f"  - 表现：{value_info['表现']}")
        if "约束" in value_info:
            self._append_constraint_info(desc_parts, value_info['约束'])
        
        if "表现形式示例" in value_info:
            desc_parts.append(f"  - 表现形式示例：{'; '.join(value_info['表现形式示例'])}")
        
        descriptions.append('\n'.join(desc_parts))
    
    def _append_constraint_info(self, desc_parts: List[str], constraint: Any) -> None:
        """
        添加约束信息到描述部分
        
        Args:
            desc_parts: 描述部分列表
            constraint: 约束信息
        """
        if isinstance(constraint, dict):
            constraint_text = []
            for k, v in constraint.items():
                if isinstance(v, dict):
                    constraint_text.append(f"{k}: {', '.join([f'{subk}: {subv}' for subk, subv in v.items()])}")
                elif isinstance(v, list):
                    constraint_text.append(f"{k}: {', '.join(v)}")
                else:
                    constraint_text.append(f"{k}: {v}")
            desc_parts.append(f"  - 约束：{'; '.join(constraint_text)}")
        else:
            desc_parts.append(f"  - 约束：{constraint}")
    
    def generate_prompt(self, patient_info: Dict[str, str], doctor_question: str, 
                       conversation_history: Optional[List[Dict[str, str]]] = None, 
                       target_personality: Optional[Dict[str, str]] = None) -> str:
        """
        生成用于指导模型生成回复的完整提示词
        
        Args:
            patient_info: 患者基本信息
            doctor_question: 当前医生的问题
            conversation_history: 对话历史
            target_personality: 目标人格特征
            
        Returns:
            完整的提示词文本
        """
        if not doctor_question:
            raise ValueError("医生问题不能为空")
            
        if not patient_info:
            logger.warning("患者信息为空，可能影响生成质量")
        
        prompt_parts = []
        
        # 任务说明
        prompt_parts.append("任务：请模拟一个具有特定人格特征的患者，根据对话历史和当前的医生话语，生成一轮符合该患者人格的自然回应。")
        prompt_parts.append("重要：请严格遵守患者的人格特征。")
        
        # 患者基本信息
        if patient_info:
            prompt_parts.append("\n患者基本信息：")
            for key, value in patient_info.items():
                prompt_parts.append(f"- {key}：{value}")
        
        # 人格定义描述
        if target_personality:
            personality_desc = self.get_personality_description(target_personality)
            prompt_parts.append(f"\n{personality_desc}")
        
        # 对话历史
        if conversation_history:
            prompt_parts.append("\n对话历史：")
            for i, turn in enumerate(conversation_history):
                prompt_parts.append(f"轮次 {i+1}：")
                if turn.get("医生"):
                    prompt_parts.append(f"  医生：{turn['医生']}")
                if turn.get("患者"):
                    prompt_parts.append(f"  患者：{turn['患者']}")
        
        # 当前医生提问
        prompt_parts.append(f"\n当前医生提问：{doctor_question}")
        
        # 关键医疗信息保留要求
        # prompt_parts.append("\n\n关键医疗信息保留要求（必须遵守）：")
        # prompt_parts.append("1. 必须保留所有症状描述：如疼痛部位、程度、持续时间等")
        # prompt_parts.append("2. 必须保留所有医疗时间信息：如发病时间、用药时长等")
        # prompt_parts.append("3. 必须保留所有药物信息：药名、剂量、用法")
        # prompt_parts.append("4. 必须保留所有检查信息：检查类型、结果、时间")
        # prompt_parts.append("5. 必须保留所有数字信息：年龄、剂量、时间数量等")
        
        # 生成要求
        prompt_parts.append("\n\n生成要求：")
        prompt_parts.append("- 轮数要求：只需要模拟患者对当前医生话语进行一轮回复，禁止再模拟医生患者自己交流。")
        prompt_parts.append("- 信息准确：如果回答中涉及患者信息，必须与患者基本信息中的事实性医疗信息保持一致")

        prompt_parts.append("- 人格体现：在回答中自然体现目标人格特征")
        prompt_parts.append("- 口语化表达：使用自然的口语化表达，避免过于书面化")
        prompt_parts.append("- 回答长度：保持合理长度，不要过长或过短")
        
        # 输出格式要求
        prompt_parts.append("\n\n输出格式要求：")
        prompt_parts.append("- 请使用'患者：'作为前缀，后跟生成的回答内容")
        prompt_parts.append("- 不要包含任何医生的提问或对话历史内容")
        prompt_parts.append("- 不要包含任何额外的说明文字或标记符号，不要包含与扮演患者回复无关的任何内容")
        prompt_parts.append("- 回答必须是完整、连贯的患者回复")
        
        final_prompt = '\n'.join(prompt_parts)
        logger.debug(f"生成的提示词长度：{len(final_prompt)}字符")
        return final_prompt
    
    # def apply_memory_constraints(self, question: str, base_answer: str, persona: Dict[str, str]) -> str:
    #     """
    #     根据病史回忆能力调整回答，避免低回忆时频繁遗忘
        
    #     Args:
    #         question: 医生的问题
    #         base_answer: 基础回答
    #         persona: 患者人格特征
            
    #     Returns:
    #         应用记忆约束后的回答
    #     """
    #     if not persona or persona.get("病史回忆能力") != "低":
    #         return base_answer  # 高回忆能力 → 保持原样
        
    #     # 检查是否为主诉或当前症状相关问题，如果是则禁止遗忘
    #     if any(keyword in question for keyword in self.MEMORY_CONSTRAINT_KEYWORDS):
    #         # 主诉问题只允许正常回答或轻微模糊
    #         if random.random() < 0.15:
    #             return self._add_vagueness_to_answer(base_answer)
    #         return base_answer
        
    #     # 根据概率应用不同程度的记忆模糊或遗忘
    #     rand = random.random()
    #     if rand < 0.08:  # 8% 完全遗忘
    #         return random.choice([
    #             "这个我真想不起来了。",
    #             "医生，这个我记不清楚了。",
    #         ])
    #     elif rand < 0.20:  # 12% 模糊/不确定
    #         return random.choice([
    #             "大概一周左右吧，也可能更久。",
    #             "好像吃了几天药，具体我不太确定。",
    #             "我记得不是很清楚，可能是上个月。"
    #         ])
    #     else:  # 80% 保持正常回答
    #         # 在正常回答里加一点点模糊修饰
    #         if random.random() < 0.15:
    #             return self._add_vagueness_to_answer(base_answer)
    #         return base_answer
    
    # def _add_vagueness_to_answer(self, answer: str) -> str:
    #     """
    #     为回答添加模糊修饰
        
    #     Args:
    #         answer: 原始回答
            
    #     Returns:
    #         添加模糊修饰后的回答
    #     """
    #     vague_suffixes = [
    #         "，不过具体我记得不是很准。",
    #         "，应该是这样吧。",
    #         "，大概是这样。"
    #     ]
    #     return answer + random.choice(vague_suffixes)

# 使用示例
if __name__ == "__main__":
    # 初始化提示词生成器
    prompt_generator = GenerateReplyPrompt()
    
    # 示例患者信息
    patient_info = {
        "姓名": "丁秀红",
        "年龄": "45岁",
        "性别": "女",
        "主诉": "胃痛一周",
        "既往史": "胃炎病史3年"
    }
    
    # 示例医生问题
    doctor_question = "丁女士，您的胃痛具体是什么样的疼痛？是胀痛、刺痛还是隐痛？"
    
    # 示例对话历史
    conversation_history = [
        {"医生": "您好，请问您哪里不舒服？", "患者": "医生您好，我胃痛，已经疼了一个星期了。"},
        {"医生": "您有胃炎病史吗？", "患者": "是的，我有胃炎病史3年了。"}
    ]
    
    # 示例目标人格
    target_personality = {
        "性格": "合作",
        "情绪状态": "焦虑",
        "病史回忆能力": "高",
        "医学理解力": "中",
        "语言表达能力": "高",
        "认知状态": "正常"
    }
    
    try:
        # 生成提示词
        prompt = prompt_generator.generate_prompt(
            patient_info=patient_info,
            doctor_question=doctor_question,
            conversation_history=conversation_history,
            target_personality=target_personality
        )
        
        print("生成的提示词：\n")
        print(prompt)
        
        # 假设模型生成了一个基础回答
        base_answer = "主要是胀痛，感觉胃里有气排不出来，尤其是吃完饭后更明显，有时候还会反酸水。"
        
        # 应用记忆约束
        final_answer = prompt_generator.apply_memory_constraints(
            question=doctor_question,
            base_answer=base_answer,
            persona=target_personality
        )
        
        print("\n\n应用约束后的最终回答：\n")
        print(final_answer)
    except Exception as e:
        logger.error(f"示例运行出错: {e}")
        print(f"错误: {e}")